import numpy as np
import matplotlib.pyplot as plt


class DoubleS:
    def __init__(self, q0: np.ndarray, q1: np.ndarray, vmax: np.ndarray, amax: np.ndarray, jmax: np.ndarray):

        self.q0 = q0
        self.q1 = q1

        self.q0_hat = 0
        self.q1_hat = 1
        self.v0_hat = 0
        self.v1_hat = 0

        self.vmax_hat = 1 / np.max(np.abs(q1 - q0) / vmax)
        self.vmin_hat = -self.vmax_hat
        self.amax_hat = 1 / np.max(np.abs(q1 - q0) / amax)
        self.amin_hat = -self.amax_hat
        self.jmax_hat = 1 / np.max(np.abs(q1 - q0) / jmax)
        self.jmin_hat = -self.jmax_hat

        # Case1 vlim == vmax
        if self.vmax_hat * self.jmax_hat >= self.amax_hat ** 2:
            self.Tj = self.amax_hat / self.jmax_hat
            self.Ta = self.Tj + self.vmax_hat / self.amax_hat
        else:
            self.Tj = np.sqrt(self.vmax_hat / self.jmax_hat)
            self.Ta = 2 * self.Tj
        self.Tv = (self.q1_hat - self.q0_hat) / self.vmax_hat - self.Ta

        # Case2 vlim < vmax
        if self.Tv < 0:
            self.Tv = 0
            if (self.q1_hat - self.q0_hat) >= 2 * self.amax_hat ** 3 / (self.jmax_hat ** 2):
                self.Tj = self.amax_hat / self.jmax_hat
                self.Ta = self.Tj / 2 + np.sqrt((self.Tj / 2) ** 2 + (self.q1_hat - self.q0_hat) / self.amax_hat)
            else:
                self.Tj = np.power((self.q1_hat - self.q0_hat) / (2 * self.jmax_hat), 1 / 3)
                self.Ta = 2 * self.Tj

        # Computation of the trajectory
        alim = self.jmax_hat * self.Tj
        self.alima = alim
        self.alimd = -alim
        self.vlim = (self.Ta - self.Tj) * alim
        self.Td = self.Ta
        self.T = self.Ta + self.Tv + self.Td

    def get_trajectory(self, t):
        # Acceleration phase
        if t < self.Tj:
            q_hat = self.q0_hat + self.v0_hat * t + self.jmax_hat * t ** 3 / 6.0
            dq_hat = self.v0_hat + self.jmax_hat * t ** 2 / 2.0
            ddq_hat = self.jmax_hat * t
            dddq_hat = self.jmax_hat
        elif t < self.Ta - self.Tj:
            q_hat = self.q0_hat + self.v0_hat * t + self.alima / 6.0 * (3.0 * t ** 2 - 3.0 * self.Tj * t + self.Tj ** 2)
            dq_hat = self.v0_hat + self.alima * (t - self.Tj / 2.0)
            ddq_hat = self.alima
            dddq_hat = 0.0
        elif t < self.Ta:
            q_hat = self.q0_hat + (self.vlim + self.v0_hat) * self.Ta / 2.0 - self.vlim * (
                    self.Ta - t) - self.jmin_hat * (self.Ta - t) ** 3 / 6.0
            dq_hat = self.vlim + self.jmin_hat * (self.Ta - t) ** 2 / 2.0
            ddq_hat = -self.jmin_hat * (self.Ta - t)
            dddq_hat = self.jmin_hat
        # Constant velocity phase
        elif t < self.Ta + self.Tv:
            q_hat = self.q0_hat + (self.vlim + self.v0_hat) * self.Ta / 2.0 + self.vlim * (t - self.Ta)
            dq_hat = self.vlim
            ddq_hat = 0.0
            dddq_hat = 0.0
        # Deceleration phase
        elif t < self.T - self.Td + self.Tj:
            q_hat = self.q1_hat - (self.vlim + self.v1_hat) * self.Td / 2.0 + self.vlim * (
                    t - self.T + self.Td) - self.jmax_hat * (t - self.T + self.Td) ** 3 / 6.0
            dq_hat = self.vlim - self.jmax_hat * (t - self.T + self.Td) ** 2 / 2.0
            ddq_hat = - self.jmax_hat * (t - self.T + self.Td)
            dddq_hat = self.jmin_hat
        elif t < self.T - self.Tj:
            q_hat = self.q1_hat - (self.vlim + self.v1_hat) * self.Td / 2.0 + self.vlim * (
                    t - self.T + self.Td) + self.alimd / 6.0 * (
                            3.0 * (t - self.T + self.Td) ** 2 - 3.0 * self.Tj * (t - self.T + self.Td) + self.Tj ** 2)
            dq_hat = self.vlim + self.alimd * (t - self.T + self.Td - self.Tj / 2.0)
            ddq_hat = self.alimd
            dddq_hat = 0.0
        elif t < self.T:
            q_hat = self.q1_hat - self.v1_hat * (self.T - t) - self.jmax_hat * (self.T - t) ** 3 / 6.0
            dq_hat = self.v1_hat + self.jmax_hat * (self.T - t) ** 2 / 2.0
            ddq_hat = -self.jmax_hat * (self.T - t)
            dddq_hat = self.jmax_hat
        else:
            q_hat = self.q1_hat
            dq_hat = self.v1_hat
            ddq_hat = 0.0
            dddq_hat = 0.0

        q = self.q0 + (self.q1 - self.q0) * q_hat
        dq = (self.q1 - self.q0) * dq_hat
        ddq = (self.q1 - self.q0) * ddq_hat
        dddq = (self.q1 - self.q0) * dddq_hat
        return q, dq, ddq, dddq

    def get_total_time(self):
        return self.T


if __name__ == "__main__":
    q0 = np.array((0.0, 0.0))
    q1 = np.array((10.0, 20.0))
    vmax = np.array((5.0, 10))
    amax = np.array((10.0, 10))
    jmax = np.array((30.0, 30))

    num = 1001

    double_s = DoubleS(q0, q1, vmax, amax, jmax)
    T = double_s.get_total_time()
    qt = []
    dqt = []
    ddqt = []
    dddqt = []

    t = np.linspace(0, np.ceil(T), num)

    for ti in t:
        trajectory = double_s.get_trajectory(ti)
        qt.append(trajectory[0])
        dqt.append(trajectory[1])
        ddqt.append(trajectory[2])
        dddqt.append(trajectory[3])

    plt.figure(1)

    plt.subplot(2, 2, 1)
    plt.plot(t, qt)
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$q (rad)$')

    plt.subplot(2, 2, 2)
    plt.plot(t, dqt)
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dot{q} (rad/s)$')

    plt.subplot(2, 2, 3)
    plt.plot(t, ddqt)
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\ddot{q} (rad/s^2)$')

    plt.subplot(2, 2, 4)
    plt.plot(t, dddqt)
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dddot{q} (rad/s^3)$')

    plt.tight_layout()
    plt.show()
